# training/experience_replay.py
import numpy as np
import random
from typing import List, Tuple
import torch

class PrioritizedExperienceReplay:
    """优先经验回放缓冲区"""
    
    def __init__(self, capacity: int = 10000, alpha: float = 0.6, beta: float = 0.4):
        self.capacity = capacity
        self.alpha = alpha  # 优先级指数
        self.beta = beta    # 重要性采样指数
        self.beta_increment = 0.001
        
        self.buffer = []
        self.priorities = np.zeros(capacity, dtype=np.float32)
        self.position = 0
        self.size = 0
    
    def add(self, experience: dict, priority: float):
        """添加经验到缓冲区"""
        if len(self.buffer) < self.capacity:
            self.buffer.append(experience)
        else:
            self.buffer[self.position] = experience
        
        # 设置优先级（避免为零）
        self.priorities[self.position] = max(priority, 1e-5) ** self.alpha
        self.position = (self.position + 1) % self.capacity
        self.size = min(self.size + 1, self.capacity)
    
    def sample(self, batch_size: int) -> Tuple[List, List, torch.Tensor]:
        """采样批次经验"""
        if self.size == 0:
            return [], [], torch.tensor([])
        
        # 计算采样概率
        priorities = self.priorities[:self.size]
        probs = priorities / priorities.sum()
        
        # 采样索引
        indices = np.random.choice(self.size, batch_size, p=probs, replace=False)
        
        # 计算重要性采样权重
        weights = (self.size * probs[indices]) ** (-self.beta)
        weights = weights / weights.max()
        weights = torch.FloatTensor(weights)
        
        # 更新beta
        self.beta = min(1.0, self.beta + self.beta_increment)
        
        # 获取经验
        batch = [self.buffer[i] for i in indices]
        
        return batch, indices, weights
    
    def update_priorities(self, indices: List[int], priorities: List[float]):
        """更新经验优先级"""
        for idx, priority in zip(indices, priorities):
            self.priorities[idx] = max(priority, 1e-5) ** self.alpha
    
    def __len__(self):
        return self.size